#!/bin/bash

# RTX 4060Ti 16G显卡显存测试脚本
# 作者：AI助手
# 功能：测试显卡显存的稳定性和完整性

# 显示彩色输出
GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
NC='\033[0m' # 无颜色

# 打印标题
echo -e "${GREEN}============================================${NC}"
echo -e "${GREEN}     RTX 4060Ti 16G 显存测试工具         ${NC}"
echo -e "${GREEN}============================================${NC}"
echo ""

# 检查是否为root用户
if [ "$EUID" -ne 0 ]; then
  echo -e "${RED}请使用root权限运行此脚本：sudo $0${NC}"
  exit 1
fi

# 确认用户意愿
echo -e "${YELLOW}此脚本将测试显卡显存的稳定性。${NC}"
echo -e "${RED}警告：显存测试会导致系统负载增加并使显卡温度升高。${NC}"
echo -e "${RED}测试期间请勿运行其他程序，以免影响测试结果。${NC}"
read -p "确定要继续吗？(y/n) " -n 1 -r
echo
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
  echo -e "${RED}操作已取消${NC}"
  exit 1
fi

# 检查NVIDIA驱动是否存在
echo -e "\n${GREEN}[步骤1] 检查NVIDIA驱动${NC}"
if ! command -v nvidia-smi &> /dev/null; then
  echo -e "${RED}未检测到NVIDIA驱动，请先安装驱动。${NC}"
  echo -e "${YELLOW}提示：可以运行 gpu_verify.sh 脚本安装驱动。${NC}"
  exit 1
fi

# 安装必要的测试工具
echo -e "\n${GREEN}[步骤2] 安装必要的测试工具${NC}"
echo "更新软件源..."
apt update

echo "安装基础工具包..."
apt install -y git build-essential cmake libx11-dev 

# 选择测试方法
echo -e "\n${GREEN}[步骤3] 选择测试方法${NC}"
echo "1. 使用CUDA示例测试显存"
echo "2. 使用GPUTest进行显存测试"
echo "3. 使用Unigine Heaven进行图形和显存测试"
echo "4. 全部测试（推荐但耗时较长）"
read -p "请选择测试方法(1-4): " test_method

# 执行CUDA示例测试
run_cuda_test() {
  echo -e "\n${YELLOW}开始CUDA显存测试...${NC}"
  
  # 安装CUDA示例
  apt install -y nvidia-cuda-samples
  
  # 检查带宽测试
  if [ -d "/usr/local/cuda/samples/1_Utilities/bandwidthTest" ]; then
    echo "编译带宽测试..."
    cd /usr/local/cuda/samples/1_Utilities/bandwidthTest
    make
    echo "运行带宽测试..."
    ./bandwidthTest
  else
    echo -e "${YELLOW}找不到CUDA示例，从GitHub下载...${NC}"
    cd /tmp
    git clone https://github.com/NVIDIA/cuda-samples.git
    cd cuda-samples/Samples/1_Utilities/bandwidthTest
    make
    echo "运行带宽测试..."
    ./bandwidthTest
  fi
  
  # 使用simpleTexture测试
  echo -e "\n${YELLOW}测试显存纹理访问...${NC}"
  if [ -d "/usr/local/cuda/samples/2_Graphics/simpleTexture" ]; then
    cd /usr/local/cuda/samples/2_Graphics/simpleTexture
    make
    ./simpleTexture
  else 
    cd /tmp/cuda-samples/Samples/2_Graphics/simpleTexture
    make
    ./simpleTexture
  fi
}

# 执行GPUTest测试
run_gputest() {
  echo -e "\n${YELLOW}开始GPUTest显存测试...${NC}"
  
  # 安装GPUTest
  apt install -y gputest
  
  # 开始显存填充测试
  echo "运行FillrateTest测试..."
  gputest --test=1 --width=1920 --height=1080 --benchmark
  
  # 开始纹理测试
  echo -e "\n运行纹理测试..."
  gputest --test=4 --width=1920 --height=1080 --benchmark
}

# 执行Unigine Heaven测试
run_heaven_test() {
  echo -e "\n${YELLOW}开始Unigine Heaven测试...${NC}"
  
  # 安装Unigine Heaven
  apt install -y unigine-heaven
  
  # 运行测试
  echo "运行Heaven基准测试..."
  heaven -benchmark -fullscreen
}

# 启动监控工具
start_monitoring() {
  echo -e "\n${GREEN}[步骤4] 启动监控${NC}"
  
  # 创建监控窗口
  gnome-terminal -- bash -c "nvidia-smi -l 1" &
  MONITOR_PID=$!
  
  # 储存PID
  echo $MONITOR_PID > /tmp/gpu_monitor_pid
  
  echo "监控已启动，测试过程中请关注温度变化..."
  sleep 3
}

# 停止监控
stop_monitoring() {
  if [ -f /tmp/gpu_monitor_pid ]; then
    MONITOR_PID=$(cat /tmp/gpu_monitor_pid)
    kill $MONITOR_PID 2>/dev/null
    rm /tmp/gpu_monitor_pid
  fi
}

# 执行测试
echo -e "\n${GREEN}[步骤5] 开始执行测试${NC}"
start_monitoring

case $test_method in
  1)
    run_cuda_test
    ;;
  2)
    run_gputest
    ;;
  3)
    run_heaven_test
    ;;
  4)
    echo -e "${YELLOW}执行全部测试...${NC}"
    run_cuda_test
    run_gputest
    run_heaven_test
    ;;
  *)
    echo -e "${RED}无效选择，退出...${NC}"
    stop_monitoring
    exit 1
    ;;
esac

# 停止监控
stop_monitoring

# 生成测试报告
echo -e "\n${GREEN}[步骤6] 生成测试报告${NC}"
REPORT_FILE="gpu_memory_test_$(date +%Y%m%d_%H%M%S).txt"

{
  echo "RTX 4060Ti 16G显存测试报告"
  echo "=========================="
  echo "测试时间: $(date)"
  echo ""
  echo "显卡信息:"
  nvidia-smi
  echo ""
  echo "显存信息:"
  nvidia-smi --query-gpu=memory.total,memory.used,memory.free --format=csv
  echo ""
  echo "测试结果分析:"
  echo "如果测试过程中没有出现图像错误、程序崩溃或系统冻结，"
  echo "且带宽测试数值接近RTX 4060Ti的理论值，则显存可能正常。"
  echo ""
  echo "RTX 4060Ti显存理论带宽：约288 GB/s"
  echo "实际测试值可能会略低于理论值，但不应显著偏离。"
} > "$REPORT_FILE"

echo -e "${GREEN}测试报告已保存至: $REPORT_FILE${NC}"

# 结论提示
echo -e "\n${GREEN}[测试完成]${NC}"
echo -e "${YELLOW}显存测试结果分析:${NC}"
echo "1. 测试过程中是否出现图像错误或程序崩溃？"
echo "   - 若是：显存可能有损坏"
echo "   - 若否：显存可能正常"
echo ""
echo "2. 带宽测试结果是否接近理论值(约288 GB/s)？"
echo "   - 若显著低于：可能是非4060Ti或有损坏"
echo "   - 若接近：可能是正常显卡"
echo ""
echo "3. 运行测试时显卡温度表现："
echo "   - 正常显卡：温度应稳定在85°C以下"
echo "   - 问题显卡：温度可能迅速升高或波动大"
echo ""
echo -e "${GREEN}感谢使用RTX 4060Ti 16G显存测试工具!${NC}" 